# Color and shape detection node for Robomaster EP
# --Students' version
# 2024/03/18
# 说明：请在以下五处Todo处按照要求完善代码，以实现期望的功能

import cv2
from cv_bridge import CvBridge, CvBridgeError
import math
import rclpy
from rclpy.node import Node
from sensor_msgs.msg import Image
from geometry_msgs.msg import Twist
import numpy as np


#定义一个寻找边缘的函数
def getContours(img,imgContour,imgColor,mask):

    contours,hierarchy = cv2.findContours(img,cv2.RETR_EXTERNAL,cv2.CHAIN_APPROX_NONE)# 寻找轮廓，第二个参数为模式，第三个为近似方法
    for cnt in contours:
        area = cv2.contourArea(cnt)
        if area > 200:                                      # 通过面积来筛选符合要求的轮廓范围，减少噪声的干扰
            cv2.drawContours(imgContour,cnt,-1,(0,255,0),3) # 画出边界
            peri = cv2.arcLength(cnt,True)                  # 计算弧长（周长）
            approx = cv2.approxPolyDP(cnt,0.02*peri,True)   # 得到角点
            objCor = len(approx)                            # 通过角点的个数来判断几何图形是几边形
            x,y,w,h = cv2.boundingRect(approx)              # 通过角点得到外接矩形
            if mask[y+h//2,x+w//2]==255:                                # 先判断边缘中心是否在mask之内
                cv2.drawContours(imgContour, cnt, -1, (0, 255, 0), 3)   # 画出边界
                '''
                # TODO:参考以下代码，补充四边形和圆形的判断。提示：当边数较多时，就认为图形为圆形！
                if objCor == 3:                                 #判断图形边数
                    objectType = "Triangle"                     #三角形
                else:
                    objectType = "None"                         #没有图形
                '''
                
                cv2.rectangle(imgColor,(x,y),(x+w,y+h),(0,255,0),2) #画出外接矩形
                
                cv2.putText(imgColor,objectType,                      #标出形状
                            (x+(w//2)-10,y+(h//2)-10),cv2.FONT_HERSHEY_COMPLEX,0.5,
                            (0,0,0),2)
                '''
                # TODO:在这里判断检测到了什么图形，如果是你需要识别的图形，则返回x,y,w,h。
                return x,y,w,h # 返回检测到的目标信息（左上角位置和宽高）
                '''
    return -1, -1, -1, -1 # 不满足条件则返回-1,-1,-1,-1


class ColorShapeDetectionNode(Node):

    def __init__(self):
        super().__init__('color_shape_detection_node')
        self.bridge = CvBridge()
        # 订阅/camera/image_color话题
        self.subscription = self.create_subscription(
            Image,
            '/camera/image_color',
            self.image_callback,
            10)
        self.subscription  # prevent unused variable warning
        '''
        # TODO:仿照人脸检测节点，在这里定义一个速度指令publisher
        self.velocity_command = Twist()
        '''
        

    # 定义回调函数
    def image_callback(self, msg):
        try:
            cv_image = self.bridge.imgmsg_to_cv2(msg, "bgr8")
            img_height, img_width, channels = cv_image.shape
        except CvBridgeError as e:
            print (e)

        imgContour = cv_image.copy()
	
        imgHSV  = cv2.cvtColor(cv_image,cv2.COLOR_BGR2HSV)

        '''
        # TODO:在这里筛选目标的颜色
        # HSV颜色范围为：[0-180, 0-255, 0-255]
        lower = np.array([5,100,100])
        upper = np.array([10,255,255])
        '''
        mask = cv2.inRange(imgHSV,lower,upper)
        mask = cv2.dilate(mask, np.ones((3, 3), np.uint8))      # 膨胀
        mask = cv2.erode(mask,np.ones((3,3),np.uint8))          # 腐蚀

        imgColor = cv2.bitwise_and(cv_image,cv_image,mask=mask)
        imgGray = cv2.cvtColor(imgColor,cv2.COLOR_BGR2GRAY)
        imgBlur = cv2.GaussianBlur(imgGray,(7,7),1)
        imgCanny = cv2.Canny(imgBlur,15,40)                     # Canny边缘检测
        imgCanny = cv2.dilate(imgCanny, np.ones((3, 3), np.uint8))
        imgCanny = cv2.erode(imgCanny, np.ones((3, 3), np.uint8))
        
        x, y, w, h =getContours(imgCanny,imgContour,imgColor,mask)

        cv2.imshow("imgColor",imgColor)
        cv2.waitKey(1)
        if(x != -1): # 如果成功检测到目标
            '''
            # 计算控制量
            # TODO:调节两个比例系数，实现目标跟随（只需调节比例系数即可，后面部分无需改动）
            self.velocity_command.linear.x = 0.0 * ((0.15 * math.sqrt(img_height ** 2 + img_width ** 2) / math.sqrt(h ** 2 + w ** 2)) - 1.0)
            self.velocity_command.angular.z = 0.0 * (0.5 - ((x + w / 2.0) / img_width))
            '''
        else:# 未检测到目标将返回-1,-1,-1,-1
            # 停车
            self.velocity_command.linear.x = 0.0
            self.velocity_command.angular.z = 0.0
        self.velocity_publisher.publish(self.velocity_command) # 发布控制指令



def main(args=None):
    try:
        rclpy.init(args=args)
        color_shape_detection_node = ColorShapeDetectionNode()
        rclpy.spin(color_shape_detection_node)
    except KeyboardInterrupt:# 关闭窗口
        print("Shutting down color shape detection node.")
        cv2.destroyAllWindows()


if __name__ == '__main__':
    main()
